package com.example.mp.loader;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.autoconfigure.vectorstore.redis.RedisVectorStoreProperties;
import org.springframework.ai.reader.JsonReader;
import org.springframework.ai.vectorstore.RedisVectorStore;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.boot.ApplicationArguments;
import org.springframework.boot.ApplicationRunner;
import org.springframework.core.io.InputStreamResource;
import org.springframework.core.io.Resource;
import org.springframework.stereotype.Component;

import java.util.Map;
import java.util.Objects;
import java.util.zip.GZIPInputStream;

/**
 * @description:
 * @Author: whopxx
 * @CreateTime: 2025-03-07
 */
@Component
public class UserHistoryBehaviorLoader implements ApplicationRunner {

    private static final Logger logger = LoggerFactory.getLogger(UserHistoryBehaviorLoader.class);

    // 定义用户历史行为数据需要读取的字段
    private static final String[] KEYS = { "custId", "behaviorType", "behaviorTime", "behaviorContent" };

    // 指向用户历史行为数据文件的路径
    @Value("classpath:/rag/user_history_behavior.json")
    private Resource data;

    private final RedisVectorStore vectorStore;

    private final RedisVectorStoreProperties properties;

    public UserHistoryBehaviorLoader(RedisVectorStore vectorStore, RedisVectorStoreProperties properties) {
        this.vectorStore = vectorStore;
        this.properties = properties;
    }

    @Override
    public void run(ApplicationArguments args) throws Exception {
        // 检查索引中的文档数量
        Map<String, Object> indexInfo = vectorStore.getJedis().ftInfo(properties.getIndex());
        logger.info("properties.getIndex: {}", properties.getIndex());
        Long numDocs = (Long) indexInfo.getOrDefault("num_docs", 0L);
        logger.info("Number of documents in the index: {}", numDocs);
        if (numDocs > 0) {
            logger.info("User history behavior embeddings already loaded. Skipping");
            return;
        }

        Resource file = data;
        // 如果文件是压缩文件，进行解压缩
        if (Objects.requireNonNull(data.getFilename()).endsWith(".gz")) {
            GZIPInputStream inputStream = new GZIPInputStream(data.getInputStream());
            file = new InputStreamResource(inputStream, "user_history_behavior.json.gz");
        }

        logger.info("Creating embeddings for user history behavior...");
        // 使用JsonReader读取用户历史行为数据
        JsonReader loader = new JsonReader(file, KEYS);
        // 将读取的数据添加到向量存储中
        vectorStore.add(loader.get());
        logger.info("Embeddings for user history behavior created.");
    }
}

